perf: reuse one StatisticsContext across ensure_distribution - #25098
perf: reuse one StatisticsContext across ensure_distribution#25098zhuqi-lucas wants to merge 3 commits into
Conversation
There was a problem hiding this comment.
🟡 Changes recommended
It introduces a breaking change to the public ensure_distribution function signature, which is re-exported and may be used by downstream crates.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR improves physical-optimizer planning performance by reusing a single StatisticsContext (and its memoization cache) across an ensure_distribution traversal, avoiding repeated recursive statistics computation on deep/wide plans.
Changes:
- Create one
StatisticsContextfor the distribution-enforcement pass and reset its cache only when a node’s plan pointer actually changes. - Thread
&StatisticsContextthroughensure_distributionand intoget_repartition_requirement_statusso child stats are memoized across the pass. - Add a deep operator-stack regression test asserting the repartition decisions remain unchanged.
File summaries
| File | Description |
|---|---|
| datafusion/physical-optimizer/src/ensure_requirements/mod.rs | Reuses a single StatisticsContext during the bottom-up distribution pass and resets cache on actual plan-pointer rewrites. |
| datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs | Threads &StatisticsContext into distribution enforcement logic to share memoized statistics across children/subtrees. |
| datafusion/core/tests/physical_optimizer/enforce_distribution.rs | Updates call sites for the new ensure_distribution signature and adds a deep-stack test for unchanged repartition behavior. |
Review details
- Files reviewed: 3/3 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| pub fn ensure_distribution( | ||
| dist_context: DistributionContext, | ||
| config: &ConfigOptions, | ||
| stats_ctx: &StatisticsContext, | ||
| ) -> Result<Transformed<DistributionContext>> { |
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #25098 +/- ##
==========================================
+ Coverage 81.80% 81.83% +0.02%
==========================================
Files 1130 1130
Lines 417842 418810 +968
Branches 417842 418810 +968
==========================================
+ Hits 341833 342742 +909
+ Misses 55877 55858 -19
- Partials 20132 20210 +78 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
a6f3f81 to
a2cb2c8
Compare
|
Hi @zhuqi-lucas, I can confirm that your use-case fits what the cache in Glad to read that you are seeing a drop in planning time of ~10% due to caching! We can probably push this a little further: instead of just passing
I have implemented something similar in #24716 for So, concretely, my suggestion is to:
This would keep the caching fix from this PR, and also allow the rule to benefit from the present and future statistics context (I will be working on #21120 as my next task). Note that there is no behavior change by default: a session without registered providers is unaffected, but it gives room for improvement without further breaking changes. I am off this week with very limited access to a computer, but I can surely offer a review or help with a PR from next week! Regarding the cache reset: your proposed solution seems safe to me, but in the future I'd like to make the cache more robust so consumers don't have to think about it. Either by introducing a unique id per constructed |
a2cb2c8 to
0b952ba
Compare
|
Thanks @asolimando — done in this PR. On the cache reset: agreed it is a footgun to leave to consumers, but since it is tangential to this PR I have filed it as a follow-up (#25141) so the robustness can be tracked separately, and kept the safe reset-on-change here. Between your two ideas I lean toward (b) holding node |
…ution get_repartition_requirement_status created a fresh StatisticsContext per child. StatisticsContext::compute recurses the child's whole subtree and carries a pointer-keyed memoization cache meant to be reused across a walk, so allocating a new context per child discards it every time and a single ensure_distribution pass recomputes shared subtree statistics O(depth) times. A fresh context also builds an empty StatisticsRegistry, so the rule never consulted registered statistics providers. Override EnsureRequirements::optimize_with_context and build one shared StatisticsContext from context.statistics_registry() per pass, threaded into ensure_distribution and get_repartition_requirement_status. This both shares the memoization cache and lets registered providers inform the distribution decision (no change by default: an empty registry behaves as before). optimize delegates to optimize_with_context via ConfigOnlyContext. The cache is keyed by raw node pointer, so it is reset after any node whose plan pointer actually changed (a rewrite can free a cached node and a later allocation could reuse its address); no-op nodes keep the cache warm. Tests: - ensure_distribution_shares_statistics_cache: counts a leaf's statistics computations under a deepening pass-through stack and asserts the shared cache saves progressively more with depth (a non-shared cache saves nothing). - ensure_distribution_uses_context_statistics_registry: a scan whose tiny real stats do not warrant round-robin is reported large by a registry provider, and the round-robin repartition then appears only via optimize_with_context. No plan changes by default; physical_optimizer and statistics suites pass. Thread the PhysicalOptimizerContext through ensure_distribution (replacing the bare ConfigOptions) so config, the statistics registry, and future CBO context are reachable without changing this public signature again. The shared StatisticsContext is built once per pass from context.statistics_registry() and threaded to get_repartition_requirement_status so the memoization cache is shared across the whole distribution walk; building it per call from the registry would lose that sharing. optimize delegates to optimize_with_context via ConfigOnlyContext. Cache is reset only when a node's plan pointer actually changed (a rewrite can free a cached node); robustness of the pointer key tracked separately in apache#25141. Tests: - ensure_distribution_shares_statistics_cache: the shared cache saves progressively more statistics recomputation as a pass-through stack deepens. - ensure_distribution_uses_context_statistics_registry: a scan reported large by a registry provider gets a round-robin repartition only via optimize_with_context, proving the registry is threaded through. No default behavior change; physical_optimizer and statistics suites pass.
0b952ba to
7389582
Compare
9bef9ca to
78c17af
Compare
The main merge pulled in apache#24809's preserve-fetch reoptimization tests, whose two ensure_distribution() calls still used the pre-refactor 2-arg (context, &config) form. Thread the PhysicalOptimizerContext and a StatisticsContext through them to match the current signature.
78c17af to
31bdfad
Compare
Which issue does this PR close?
None filed; small self-contained perf fix. Rationale below.
Rationale for this change
get_repartition_requirement_statuscreates a freshStatisticsContext::new()once per child.StatisticsContext::computerecurses the child's whole subtree and carries a pointer-keyed memoization cache its own docstring describes as a "per-call memoization cache" meant to be reused across a walk. Allocating a new context per child discards that cache every time, so a singleensure_distributionpass recomputes shared subtree statisticsO(depth)times.On a deep/wide plan this is measurable. In our deployment (
EnsureRequirementsruns several times over a ~200-node plan) sharing the cache cut physical planning by ~10% with no plan change.What changes are included in this PR?
StatisticsContextthrough theensure_distributiontransform_up(pass&StatisticsContextintoget_repartition_requirement_status) so each subtree's statistics are computed once per pass.StatsCacheis keyed by raw node pointer, andensure_distributionreturnsTransformed::yesunconditionally, so the cache reset is keyed on whether the node's plan pointer actually changed (Arc::ptr_eqbefore/after). A node that changed may have freed a cached child (which would make a stale pointer key unsafe); a node that made no change cannot, so the cache safely persists across the no-op nodes that dominate a deep plan.(A second per-child
StatisticsContext::new()inPlanSize::from_plan/enforce_distribution_relationshipscan get the same treatment; left as a follow-up to keep this PR focused.)Are these changes tested?
Yes. New test
ensure_distribution_shares_statistics_cacheputs a leaf that counts its own statistics computations under a stack of pass-through operators, runs the distribution pass with a shared context vs a fresh-per-node context, and asserts the shared cache saves progressively more as the stack deepens. A cache that is not actually shared (e.g. reset on every node) saves nothing and fails the test — which a plan-output assertion cannot catch, since the optimized plan is identical either way.Existing suites remain green and unchanged:
datafusion --test core_integration physical_optimizer(569 passed) anddatafusion-physical-planstatistics tests (96 passed).Are there any user-facing changes?
No. Internal physical-optimizer performance only; planner output is identical.